"""
Coze OpenAPI客户端实现
支持文件上传和多模态对话功能
使用官方 cozepy SDK
"""

import asyncio
import secrets
import tempfile
import os
import json
import logging
from typing import Optional, Tuple, Dict, Any, List
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

from cozepy import (
    COZE_CN_BASE_URL,
    Coze,
    TokenAuth,
    Message,
    MessageObjectString,
    ChatEventType,
)

logger = logging.getLogger(__name__)


@dataclass
class CozeAPIConfig:
    """Coze API配置"""
    api_key: str
    bot_id: str
    base_url: str = "https://api.coze.cn"
    timeout: int = 30
    max_retries: int = 3
    retry_delay: float = 1.0


@dataclass
class FileUploadResult:
    """文件上传结果"""
    file_id: str
    filename: str
    file_type: str
    file_size: int
    upload_time: datetime


@dataclass
class ChatResponse:
    """对话响应"""
    response_text: str
    confidence_score: float
    processing_time_ms: int
    model_used: str
    timestamp: datetime


class CozeAPIError(Exception):
    """Coze API错误基类"""
    pass


class CozeAuthenticationError(CozeAPIError):
    """认证错误"""
    pass


class CozeRateLimitError(CozeAPIError):
    """频率限制错误"""
    pass


class CozeNetworkError(CozeAPIError):
    """网络错误"""
    pass


class CozeOpenAPIClient:
    """
    Coze OpenAPI客户端
    支持文件上传和多模态对话功能
    使用官方 cozepy SDK
    """
    
    def __init__(self, config: CozeAPIConfig):
        """
        初始化Coze OpenAPI客户端
        
        Args:
            config: Coze API配置
        """
        self.config = config
        self.coze_client: Optional[Coze] = None
        
    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self.connect()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.disconnect()
        
    async def connect(self):
        """建立连接"""
        if self.coze_client is None:
            self.coze_client = Coze(
                auth=TokenAuth(self.config.api_key),
                base_url=self.config.base_url
            )
            logger.info("Coze OpenAPI客户端连接已建立")
            
    async def disconnect(self):
        """关闭连接"""
        if self.coze_client:
            self.coze_client = None
            logger.info("Coze OpenAPI客户端连接已关闭")
            

                
    async def upload_file(self, file_path: str, file_type: str = "auto") -> FileUploadResult:
        """
        上传文件到Coze平台
        
        Args:
            file_path: 文件路径
            file_type: 文件类型 ("image", "audio", "auto")
            
        Returns:
            文件上传结果
            
        Raises:
            CozeAPIError: 上传失败
        """
        if not self.coze_client:
            await self.connect()
            
        if not os.path.exists(file_path):
            raise CozeAPIError(f"文件不存在: {file_path}")
            
        file_path = Path(file_path)
        file_size = file_path.stat().st_size
        
        # 自动检测文件类型
        if file_type == "auto":
            suffix = file_path.suffix.lower()
            if suffix in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
                file_type = "image"
            elif suffix in ['.mp3', '.wav', '.ogg', '.m4a', '.aac']:
                file_type = "audio"
            else:
                file_type = "file"
                
        try:
            # 使用 SDK 上传文件
            file_result = self.coze_client.files.upload(file=str(file_path))
            
            return FileUploadResult(
                file_id=file_result.id,
                filename=file_result.file_name,
                file_type=file_type,
                file_size=file_result.bytes,
                upload_time=datetime.now()
            )
            
        except Exception as e:
            logger.error(f"文件上传失败: {e}")
            raise CozeAPIError(f"文件上传失败: {e}")
            

        
    async def upload_files(self, image_path: str, audio_path: Optional[str] = None) -> Tuple[str, Optional[str]]:
        """
        上传图片和音频文件
        
        Args:
            image_path: 图片文件路径
            audio_path: 音频文件路径（可选）
            
        Returns:
            (image_file_id, audio_file_id) 元组
            
        Raises:
            CozeAPIError: 上传失败
        """
        try:
            # 上传图片文件
            image_result = await self.upload_file(image_path, "image")
            image_file_id = image_result.file_id
            
            # 上传音频文件（如果提供）
            audio_file_id = None
            if audio_path and os.path.exists(audio_path):
                audio_result = await self.upload_file(audio_path, "audio")
                audio_file_id = audio_result.file_id
                
            logger.info(f"文件上传成功 - 图片ID: {image_file_id}, 音频ID: {audio_file_id}")
            return image_file_id, audio_file_id
            
        except Exception as e:
            logger.error(f"批量文件上传失败: {e}")
            raise CozeAPIError(f"批量文件上传失败: {e}")
            
    async def chat_with_vision(
        self, 
        image_file_id: str, 
        audio_file_id: Optional[str] = None,
        query: str = "请描述这张图片的内容",
        user_id: Optional[str] = None
    ) -> ChatResponse:
        """
        发起多模态对话
        
        Args:
            image_file_id: 图片文件ID
            audio_file_id: 音频文件ID（可选）
            query: 查询文本
            user_id: 用户ID（可选）
            
        Returns:
            对话响应
            
        Raises:
            CozeAPIError: 对话失败
        """
        if not self.coze_client:
            await self.connect()
            
        if not user_id:
            user_id = secrets.token_urlsafe(16)
            
        start_time = datetime.now()
        
        try:
            # 构建消息对象列表
            message_objects = [
                MessageObjectString.build_image(file_id=image_file_id)
            ]
            
            # 添加音频对象（如果有）
            if audio_file_id:
                message_objects.append(
                    MessageObjectString.build_audio(file_id=audio_file_id)
                )
                
            # 添加文本查询
            if query:
                message_objects.append(
                    MessageObjectString.build_text(text=query)
                )
            
            # 发起流式对话
            stream = self.coze_client.chat.stream(
                bot_id=self.config.bot_id,
                user_id=user_id,
                additional_messages=[
                    Message.build_user_question_objects(message_objects)
                ],
            )
            
            # 收集响应文本
            response_text = ""
            for event in stream:
                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                    response_text += event.message.content
                    
            if not response_text:
                response_text = "抱歉，我无法理解这个内容。"
                
            processing_time = int((datetime.now() - start_time).total_seconds() * 1000)
            
            return ChatResponse(
                response_text=response_text,
                confidence_score=0.8,  # 默认置信度
                processing_time_ms=processing_time,
                model_used="coze",
                timestamp=datetime.now()
            )
            
        except Exception as e:
            logger.error(f"多模态对话失败: {e}")
            raise CozeAPIError(f"多模态对话失败: {e}")
            
    def save_image_to_temp(self, image_data: bytes, format: str = "JPEG") -> str:
        """
        将图像数据保存到临时文件（同步方法）
        
        Args:
            image_data: 图像二进制数据
            format: 图像格式
            
        Returns:
            临时文件路径
        """
        suffix = ".jpg" if format.upper() == "JPEG" else f".{format.lower()}"
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_file.write(image_data)
            temp_path = temp_file.name
            
        logger.debug(f"图像数据已保存到临时文件: {temp_path}")
        return temp_path
        
    def save_audio_to_temp(self, audio_data: bytes, format: str = "WAV") -> str:
        """
        将音频数据保存到临时文件（同步方法）
        
        Args:
            audio_data: 音频二进制数据
            format: 音频格式
            
        Returns:
            临时文件路径
        """
        suffix = f".{format.lower()}"
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_file.write(audio_data)
            temp_path = temp_file.name
            
        logger.debug(f"音频数据已保存到临时文件: {temp_path}")
        return temp_path
        
    def cleanup_temp_file(self, file_path: str):
        """
        清理临时文件
        
        Args:
            file_path: 文件路径
        """
        try:
            if os.path.exists(file_path):
                os.unlink(file_path)
                logger.debug(f"临时文件已清理: {file_path}")
        except Exception as e:
            logger.warning(f"清理临时文件失败: {e}")
            
    async def process_vision_query(
        self, 
        image_data: bytes, 
        audio_data: Optional[bytes] = None,
        query: str = "请描述这张图片的内容"
    ) -> ChatResponse:
        """
        处理视觉查询的完整流程
        
        Args:
            image_data: 图像二进制数据
            audio_data: 音频二进制数据（可选）
            query: 查询文本
            
        Returns:
            对话响应
        """
        image_temp_path = None
        audio_temp_path = None
        
        try:
            # 保存图像到临时文件
            image_temp_path = self.save_image_to_temp(image_data)
            
            # 保存音频到临时文件（如果有）
            if audio_data:
                audio_temp_path = self.save_audio_to_temp(audio_data)
                
            # 上传文件
            image_file_id, audio_file_id = await self.upload_files(
                image_temp_path, 
                audio_temp_path
            )
            
            # 发起对话
            response = await self.chat_with_vision(
                image_file_id, 
                audio_file_id, 
                query
            )
            
            return response
            
        finally:
            # 清理临时文件
            if image_temp_path:
                self.cleanup_temp_file(image_temp_path)
            if audio_temp_path:
                self.cleanup_temp_file(audio_temp_path)